package net.cyndeline.scalarlib.rlgraph.util

import scalax.collection.GraphEdge.UnDiEdge
import scalax.collection.immutable.Graph
import scalax.collection.GraphPredef._
import scala.reflect.runtime.universe.TypeTag
import scala.collection.mutable

/**
 * Converts a graph with edge classes that extend the undirected edge to a graph that uses undirected edges instead.
 * @param nodeFactory Converts nodes to a different type, if specified by the user. Defaults to a factory that returns
 *                    its input.
 */
class UndirectedGraphTypeConverter[VType, NodeType : TypeTag](nodeFactory: VType => NodeType) extends GraphTypeConverter[VType, NodeType] {

  /**
   * @param graph Graph to convert.
   * @tparam EType Type of edge used in graph.
   * @return A graph with the same topology as the input, but with nodes created by the factory, and
   *         edges being undirected.
   */
  def convertToUndirected[EType[X] <: UnDiEdge[X]: ({type l[M[_]] = TypeTag[M[NodeType]]})#l]
    (graph: Graph[VType, EType]): Graph[NodeType, UnDiEdge] = {
    val addedNodes = new mutable.HashMap[VType, NodeType]()

    val nodes: Vector[NodeType] = (for {
      n <- graph.nodes
      outer: VType = n
      converted: NodeType = convertNode(outer, addedNodes)
    } yield converted).toVector

    val edges: Vector[UnDiEdge[NodeType]] = (for {
      edge <- graph.edges
      outerEdge = edge.toOuter
      from = addedNodes.get(outerEdge._1).getOrElse {
        convertNode(outerEdge._1, addedNodes)
      }
      to = addedNodes.get(outerEdge._2).getOrElse {
        convertNode(outerEdge._2, addedNodes)
      }
      unDiEdge: UnDiEdge[NodeType] = from~to
    } yield unDiEdge).toVector

    Graph.from(nodes, edges)
  }

  private def convertNode(v: VType, mapping: mutable.HashMap[VType, NodeType]): NodeType = {
    val converted = nodeFactory(v)
    mapping += v -> converted
    converted
  }

}

/**
 * Used for injection.
 */
trait GraphTypeConverter[VType, NodeType] {
  def convertToUndirected[EType[X] <: UnDiEdge[X]: ({type l[M[_]] = TypeTag[M[NodeType]]})#l](graph: Graph[VType, EType]): Graph[NodeType, UnDiEdge]
}

/**
 * Helper object available to the user.
 */
object UndirectedGraphTypeConverter {

  /**
   * @param graph Graph to convert.
   * @tparam VType Vertex type used in the graph.
   * @tparam EType Edge type used in the graph.
   * @return A graph with every edge converted to undirected edges.
   */
  def apply[VType : TypeTag, EType[X] <: UnDiEdge[X]: ({type l[M[_]] = TypeTag[M[VType]]})#l](graph: Graph[VType, EType]): Graph[VType, UnDiEdge]
    = new UndirectedGraphTypeConverter[VType, VType]((v: VType) => v).convertToUndirected(graph)

  /**
   * Converts a graph using a specified node factory.
   * @param graph Graph to convert.
   * @param nodeFactory A function that takes a vertex from the original graph and outputs another object. Will be
   *                    called once for every vertex in the graph, and only once per vertex.
   * @tparam VType Vertex type used in the graph.
   * @tparam NodeType Type that vertices in the original graph should be converted to.
   * @tparam EType Edge type used in the graph.
   * @return A graph with undirected edges, and every vertex converted by the factory.
   */
  def apply[VType : TypeTag, NodeType : TypeTag, EType[X] <: UnDiEdge[X]: ({type l[M[_]] = TypeTag[M[NodeType]]})#l]
    (graph: Graph[VType, EType], nodeFactory: VType => NodeType): Graph[NodeType, UnDiEdge]
    = new UndirectedGraphTypeConverter[VType, NodeType](nodeFactory).convertToUndirected(graph)

}
